package com.lrj.datastream.source;

import com.lrj.datastream.domain.ClickEvent;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.util.Random;
import java.util.concurrent.TimeUnit;

/**
 * 自定义数据源
 *
 * @author lrj
 * @date 2022/3/28 14:21
 */
public class TestCustomSource {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        // 默认是单并行度的,parallelism只能是1,超过1会报错
        DataStreamSource<ClickEvent> customStream = env.addSource(new ClickSource()).setParallelism(1);
        customStream.print("SourceFunction");

        // 并行的需要实现 ParallelSourceFunction
        DataStreamSource<ClickEvent> customStream2 = env.addSource(new ParallelSource()).setParallelism(2);
        customStream2.print("ParallelSourceFunction");
        env.execute(TestCustomSource.class.getName());
    }

    private static class ClickSource implements SourceFunction<ClickEvent> {

        private volatile boolean running = true;

        @Override
        public void run(SourceContext<ClickEvent> ctx) throws Exception {
            Random random = new Random();
            String[] users = {"Mary", "Bob", "Alice", "Jack"};
            String[] urls = {"./home", "./product?id=%s", "./cart"};
            while (running) {
                String user = users[random.nextInt(users.length)];
                String url = String.format(urls[random.nextInt(urls.length)], random.nextInt(100));
                ClickEvent clickEvent = new ClickEvent(user, url, System.currentTimeMillis());
                ctx.collect(clickEvent);
                TimeUnit.SECONDS.sleep(1);
            }
        }

        @Override
        public void cancel() {
            running = false;
        }
    }

    private static class ParallelSource implements ParallelSourceFunction<ClickEvent> {

        @Override
        public void run(SourceContext<ClickEvent> ctx) throws Exception {
            Random random = new Random();
            String[] users = {"Mary", "Bob", "Alice", "Jack"};
            String[] urls = {"./home", "./product?id=%s", "./cart"};
            int userLen = users.length;
            int urlLen = random.nextInt(urls.length);
            while (!Thread.currentThread().isInterrupted()) {
                String user = users[random.nextInt(userLen)];
                String url = String.format(urls[urlLen], random.nextInt(100));
                ClickEvent clickEvent = new ClickEvent(user, url, System.currentTimeMillis());
                ctx.collect(clickEvent);
                TimeUnit.SECONDS.sleep(1);
            }
        }

        @Override
        public void cancel() {
            try {
                Thread.currentThread().interrupt();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

}
